{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "94508a4f-ec1a-4a85-b104-63b6dd18a791",
   "metadata": {},
   "source": [
    "# Chain of Table Notebook\n",
    "\n",
    "<a href=\"https://colab.research.google.com/github/run-llama/llama-hub/blob/main/llama_hub/llama_packs/tables/chain_of_table/chain_of_table.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n",
    "\n",
    "In this notebook we highlight our implementation of the [\"Chain of Table\" paper by Wang et al.](https://arxiv.org/pdf/2401.04398v1.pdf).\n",
    "\n",
    "Chain-of-Table proposes the following: given a user query over tabular data, plan out a sequence of tabular operations over the table to retrieve the right information in order to satisfy the user query. The updated table is explicitly used/modified throughout the intermediate chain (unlike chain-of-thought/ReAct which uses generic thoughts). \n",
    "\n",
    "There is a fixed set of tabular operations that are defined in the paper:\n",
    "- `f_add_column`\n",
    "- `f_select_row`\n",
    "- `f_select_column`\n",
    "- `f_group_by`\n",
    "- `f_sort_by`\n",
    "\n",
    "We implemented the paper based on the prompts described in the paper, and adapted it to get it working. That said, this is marked as beta, so there may still be kinks to work through. Do you have suggestions / contributions on how to improve the robustness? Let us know! "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0646bc55-493e-4836-9df3-0f42814bc4fd",
   "metadata": {},
   "source": [
    "## Download Data\n",
    "\n",
    "We use the [WikiTableQuestions dataset](https://ppasupat.github.io/WikiTableQuestions/) (Pasupat and Liang 2015) as our test dataset.\n",
    "\n",
    "WikiTableQuestions is a question-answering dataset over various semi-structured tables taken from Wikipedia. These tables range in size from a few rows/columns to mnay rows. Some columns may contain multi-part information as well (e.g. a temperature column may contain both Fahrenheight and Celsius)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e52689a4-152b-49ef-8bf1-f588a692c5ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget \"https://github.com/ppasupat/WikiTableQuestions/releases/download/v1.0.2/WikiTableQuestions-1.0.2-compact.zip\" -O data.zip\n",
    "!unzip data.zip"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "933d0434-e8d5-45ab-9ed0-c8b192fa6d44",
   "metadata": {},
   "source": [
    "Let's visualize a sample df."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f243b923-b8f8-43d3-b469-a2fbbbb64b37",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Year</th>\n",
       "      <th>Winner</th>\n",
       "      <th>Jockey</th>\n",
       "      <th>Trainer</th>\n",
       "      <th>Owner</th>\n",
       "      <th>Breeder</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1919</td>\n",
       "      <td>Sir Barton</td>\n",
       "      <td>Johnny Loftus</td>\n",
       "      <td>H. Guy Bedwell</td>\n",
       "      <td>J. K. L. Ross</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1930</td>\n",
       "      <td>Gallant Fox</td>\n",
       "      <td>Earl Sande</td>\n",
       "      <td>Jim Fitzsimmons</td>\n",
       "      <td>Belair Stud</td>\n",
       "      <td>Belair Stud</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1935</td>\n",
       "      <td>Omaha</td>\n",
       "      <td>Willie Saunders</td>\n",
       "      <td>Jim Fitzsimmons</td>\n",
       "      <td>Belair Stud</td>\n",
       "      <td>Belair Stud</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>1937</td>\n",
       "      <td>War Admiral</td>\n",
       "      <td>Charley Kurtsinger</td>\n",
       "      <td>George H. Conway</td>\n",
       "      <td>Samuel D. Riddle</td>\n",
       "      <td>Samuel D. Riddle</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>1941</td>\n",
       "      <td>Whirlaway</td>\n",
       "      <td>Eddie Arcaro</td>\n",
       "      <td>Ben A. Jones</td>\n",
       "      <td>Calumet Farm</td>\n",
       "      <td>Calumet Farm</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>1943</td>\n",
       "      <td>Count Fleet</td>\n",
       "      <td>Johnny Longden</td>\n",
       "      <td>Don Cameron</td>\n",
       "      <td>Fannie Hertz</td>\n",
       "      <td>Fannie Hertz</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>1946</td>\n",
       "      <td>Assault</td>\n",
       "      <td>Warren Mehrtens</td>\n",
       "      <td>Max Hirsch</td>\n",
       "      <td>King Ranch</td>\n",
       "      <td>King Ranch</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>1948</td>\n",
       "      <td>Citation</td>\n",
       "      <td>Eddie Arcaro</td>\n",
       "      <td>Horace A. Jones</td>\n",
       "      <td>Calumet Farm</td>\n",
       "      <td>Calumet Farm</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>1973</td>\n",
       "      <td>Secretariat</td>\n",
       "      <td>Ron Turcotte</td>\n",
       "      <td>Lucien Laurin</td>\n",
       "      <td>Meadow Stable</td>\n",
       "      <td>Meadow Stud</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>1977</td>\n",
       "      <td>Seattle Slew</td>\n",
       "      <td>Jean Cruguet</td>\n",
       "      <td>William H. Turner, Jr.</td>\n",
       "      <td>Karen L. Taylor</td>\n",
       "      <td>Ben S. Castleman</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>1978</td>\n",
       "      <td>Affirmed</td>\n",
       "      <td>Steve Cauthen</td>\n",
       "      <td>Laz Barrera</td>\n",
       "      <td>Harbor View Farm</td>\n",
       "      <td>Harbor View Farm</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    Year        Winner              Jockey                 Trainer  \\\n",
       "0   1919    Sir Barton       Johnny Loftus          H. Guy Bedwell   \n",
       "1   1930   Gallant Fox          Earl Sande         Jim Fitzsimmons   \n",
       "2   1935         Omaha     Willie Saunders         Jim Fitzsimmons   \n",
       "3   1937   War Admiral  Charley Kurtsinger        George H. Conway   \n",
       "4   1941     Whirlaway        Eddie Arcaro            Ben A. Jones   \n",
       "5   1943   Count Fleet      Johnny Longden             Don Cameron   \n",
       "6   1946       Assault     Warren Mehrtens              Max Hirsch   \n",
       "7   1948      Citation        Eddie Arcaro         Horace A. Jones   \n",
       "8   1973   Secretariat        Ron Turcotte           Lucien Laurin   \n",
       "9   1977  Seattle Slew        Jean Cruguet  William H. Turner, Jr.   \n",
       "10  1978      Affirmed       Steve Cauthen             Laz Barrera   \n",
       "\n",
       "               Owner           Breeder  \n",
       "0      J. K. L. Ross               NaN  \n",
       "1        Belair Stud       Belair Stud  \n",
       "2        Belair Stud       Belair Stud  \n",
       "3   Samuel D. Riddle  Samuel D. Riddle  \n",
       "4       Calumet Farm      Calumet Farm  \n",
       "5       Fannie Hertz      Fannie Hertz  \n",
       "6         King Ranch        King Ranch  \n",
       "7       Calumet Farm      Calumet Farm  \n",
       "8      Meadow Stable       Meadow Stud  \n",
       "9    Karen L. Taylor  Ben S. Castleman  \n",
       "10  Harbor View Farm  Harbor View Farm  "
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "\n",
    "df = pd.read_csv(\"./WikiTableQuestions/csv/200-csv/3.csv\")\n",
    "df"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed9206f7-8271-4436-8b77-c653e8f1620f",
   "metadata": {},
   "source": [
    "## Load Pack / Setup\n",
    "\n",
    "Now we do `download_llama_pack` to load the Chain-Of-Table LlamaPack (you can also import the module directly if using the llama-hub package).\n",
    "\n",
    "We will also optionally setup observability/tracing so we can observe the intermediate steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "74dfb6f1-68c9-4378-83c0-318e12b1cb41",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Option: if developing with the llama_hub package\n",
    "# from llama_hub.llama_packs.tables.chain_of_table.base import (\n",
    "#     ChainOfTableQueryEngine,\n",
    "#     serialize_table\n",
    "# )\n",
    "\n",
    "# Option: download llama_pack\n",
    "from llama_index.llama_pack import download_llama_pack\n",
    "\n",
    "download_llama_pack(\n",
    "    \"ChainOfTablePack\",\n",
    "    \"./chain_of_table_pack\",\n",
    "    skip_load=True,\n",
    "    # leave the below line commented out if using the notebook on main\n",
    "    # llama_hub_url=\"https://raw.githubusercontent.com/run-llama/llama-hub/jerry/add_chain_of_table/llama_hub\"\n",
    ")\n",
    "from chain_of_table_pack.base import ChainOfTableQueryEngine, serialize_table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e48f36b3-9278-4da9-891c-3af4de747528",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms import OpenAI\n",
    "\n",
    "llm = OpenAI(model=\"gpt-4-1106-preview\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "175f9ccc-1d6c-4fe6-b8ac-7c416cd39f7b",
   "metadata": {},
   "source": [
    "### Optional: Setup Observability\n",
    "\n",
    "Here we will use our Arize Phoenix integration to view traces through the query engine."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7b2dbe01-cabd-449f-bc13-25a0b48bb5d6",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/jerryliu/Programming/llama-hub/.venv/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "🌍 To view the Phoenix app in your browser, visit http://127.0.0.1:6006/\n",
      "📺 To view the Phoenix app in a notebook, run `px.active_session().view()`\n",
      "📖 For more information on how to use Phoenix, check out https://docs.arize.com/phoenix\n"
     ]
    }
   ],
   "source": [
    "import phoenix as px\n",
    "import llama_index\n",
    "\n",
    "px.launch_app()\n",
    "llama_index.set_global_handler(\"arize_phoenix\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40213404-e3a2-4412-a217-82fa95f06cec",
   "metadata": {},
   "source": [
    "## Try out some Queries\n",
    "\n",
    "Now let's try out our `ChainOfTableQueryEngine`!\n",
    "\n",
    "We run it over a few different tables."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dcc5cc4b-793a-4bae-bb75-9fd045acc282",
   "metadata": {},
   "source": [
    "### Example: Movie Awards Table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "cec53d01-ac5f-484d-a7cf-8dcc41c30bc9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "df = pd.read_csv(\"~/Downloads/WikiTableQuestions/csv/200-csv/11.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "fc862176-90d3-4c27-87ef-533fe9870c49",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Award</th>\n",
       "      <th>Category</th>\n",
       "      <th>Nominee</th>\n",
       "      <th>Result</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Academy Awards, 1972</td>\n",
       "      <td>Best Picture</td>\n",
       "      <td>Phillip D'Antoni</td>\n",
       "      <td>Won</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Academy Awards, 1972</td>\n",
       "      <td>Best Director</td>\n",
       "      <td>William Friedkin</td>\n",
       "      <td>Won</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Academy Awards, 1972</td>\n",
       "      <td>Best Actor</td>\n",
       "      <td>Gene Hackman</td>\n",
       "      <td>Won</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Academy Awards, 1972</td>\n",
       "      <td>Best Adapted Screenplay</td>\n",
       "      <td>Ernest Tidyman</td>\n",
       "      <td>Won</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Academy Awards, 1972</td>\n",
       "      <td>Film Editing</td>\n",
       "      <td>Gerald B. Greenberg</td>\n",
       "      <td>Won</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>Academy Awards, 1972</td>\n",
       "      <td>Best Supporting Actor</td>\n",
       "      <td>Roy Scheider</td>\n",
       "      <td>Nominated</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>Academy Awards, 1972</td>\n",
       "      <td>Best Cinematography</td>\n",
       "      <td>Owen Roizman</td>\n",
       "      <td>Nominated</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>Academy Awards, 1972</td>\n",
       "      <td>Best Sound</td>\n",
       "      <td>Theodore Soderberg\\nChristopher Newman</td>\n",
       "      <td>Nominated</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>American Cinema Editors, 1972</td>\n",
       "      <td>Best Edited Feature Film</td>\n",
       "      <td>Gerald B. Greenberg</td>\n",
       "      <td>Nominated</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>BAFTA, 1972</td>\n",
       "      <td>Best Actor</td>\n",
       "      <td>Gene Hackman</td>\n",
       "      <td>Won</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>BAFTA, 1972</td>\n",
       "      <td>Best Film Editing</td>\n",
       "      <td>Gerald B. Greenberg</td>\n",
       "      <td>Won</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>BAFTA, 1972</td>\n",
       "      <td>Best Direction</td>\n",
       "      <td>William Friedkin</td>\n",
       "      <td>Nominated</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>BAFTA, 1972</td>\n",
       "      <td>Best Film</td>\n",
       "      <td>Philip D'Antoni</td>\n",
       "      <td>Nominated</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13</th>\n",
       "      <td>BAFTA, 1972</td>\n",
       "      <td>Best Sound Track</td>\n",
       "      <td>Christopher Newman\\nTheodore Soderberg</td>\n",
       "      <td>Nominated</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14</th>\n",
       "      <td>David di Donatello Award, 1972</td>\n",
       "      <td>Best Foreign Film</td>\n",
       "      <td>Philip D'Antoni</td>\n",
       "      <td>Won</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>15</th>\n",
       "      <td>Directors Guild of America, 1972</td>\n",
       "      <td>Outstanding Directorial Achievement</td>\n",
       "      <td>William Friedkin</td>\n",
       "      <td>Won</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16</th>\n",
       "      <td>Edgar Allan Poe Awards, 1972</td>\n",
       "      <td>Best Motion Picture</td>\n",
       "      <td>Ernest Tidyman</td>\n",
       "      <td>Won</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>17</th>\n",
       "      <td>Golden Globe Awards, 1972</td>\n",
       "      <td>Best Motion Picture</td>\n",
       "      <td>Phillip D'Antoni</td>\n",
       "      <td>Won</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>18</th>\n",
       "      <td>Golden Globe Awards, 1972</td>\n",
       "      <td>Best Director</td>\n",
       "      <td>William Friedkin</td>\n",
       "      <td>Won</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>19</th>\n",
       "      <td>Golden Globe Awards, 1972</td>\n",
       "      <td>Best Actor</td>\n",
       "      <td>Gene Hackman</td>\n",
       "      <td>Won</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>20</th>\n",
       "      <td>Golden Globe Awards, 1972</td>\n",
       "      <td>Best Screenplay</td>\n",
       "      <td>Ernest Tidyman</td>\n",
       "      <td>Nominated</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>21</th>\n",
       "      <td>Kansas City Film Critics Circle, 1972</td>\n",
       "      <td>Best Actor</td>\n",
       "      <td>Gene Hackman</td>\n",
       "      <td>Won</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>22</th>\n",
       "      <td>Kansas City Film Critics Circle, 1972</td>\n",
       "      <td>Best Film</td>\n",
       "      <td>Ernest Tidyman</td>\n",
       "      <td>Won</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>23</th>\n",
       "      <td>National Society of Film Critics, 1972</td>\n",
       "      <td>Best Actor</td>\n",
       "      <td>Gene Hackman</td>\n",
       "      <td>Nominated</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>24</th>\n",
       "      <td>New York Film Critics Circle, 1971</td>\n",
       "      <td>Best Actor</td>\n",
       "      <td>Gene Hackman</td>\n",
       "      <td>Won</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>25</th>\n",
       "      <td>New York Film Critics Circle, 1971</td>\n",
       "      <td>Best Film</td>\n",
       "      <td>Ernest Tidyman</td>\n",
       "      <td>Nominated</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>26</th>\n",
       "      <td>Writers Guild of America, 1972</td>\n",
       "      <td>Best Drama Adaptation</td>\n",
       "      <td>Ernest Tidyman</td>\n",
       "      <td>Nominated</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                     Award  \\\n",
       "0                     Academy Awards, 1972   \n",
       "1                     Academy Awards, 1972   \n",
       "2                     Academy Awards, 1972   \n",
       "3                     Academy Awards, 1972   \n",
       "4                     Academy Awards, 1972   \n",
       "5                     Academy Awards, 1972   \n",
       "6                     Academy Awards, 1972   \n",
       "7                     Academy Awards, 1972   \n",
       "8            American Cinema Editors, 1972   \n",
       "9                              BAFTA, 1972   \n",
       "10                             BAFTA, 1972   \n",
       "11                             BAFTA, 1972   \n",
       "12                             BAFTA, 1972   \n",
       "13                             BAFTA, 1972   \n",
       "14          David di Donatello Award, 1972   \n",
       "15        Directors Guild of America, 1972   \n",
       "16            Edgar Allan Poe Awards, 1972   \n",
       "17               Golden Globe Awards, 1972   \n",
       "18               Golden Globe Awards, 1972   \n",
       "19               Golden Globe Awards, 1972   \n",
       "20               Golden Globe Awards, 1972   \n",
       "21   Kansas City Film Critics Circle, 1972   \n",
       "22   Kansas City Film Critics Circle, 1972   \n",
       "23  National Society of Film Critics, 1972   \n",
       "24      New York Film Critics Circle, 1971   \n",
       "25      New York Film Critics Circle, 1971   \n",
       "26          Writers Guild of America, 1972   \n",
       "\n",
       "                               Category  \\\n",
       "0                          Best Picture   \n",
       "1                         Best Director   \n",
       "2                            Best Actor   \n",
       "3               Best Adapted Screenplay   \n",
       "4                          Film Editing   \n",
       "5                 Best Supporting Actor   \n",
       "6                   Best Cinematography   \n",
       "7                            Best Sound   \n",
       "8              Best Edited Feature Film   \n",
       "9                            Best Actor   \n",
       "10                    Best Film Editing   \n",
       "11                       Best Direction   \n",
       "12                            Best Film   \n",
       "13                     Best Sound Track   \n",
       "14                    Best Foreign Film   \n",
       "15  Outstanding Directorial Achievement   \n",
       "16                  Best Motion Picture   \n",
       "17                  Best Motion Picture   \n",
       "18                        Best Director   \n",
       "19                           Best Actor   \n",
       "20                      Best Screenplay   \n",
       "21                           Best Actor   \n",
       "22                            Best Film   \n",
       "23                           Best Actor   \n",
       "24                           Best Actor   \n",
       "25                            Best Film   \n",
       "26                Best Drama Adaptation   \n",
       "\n",
       "                                   Nominee     Result  \n",
       "0                         Phillip D'Antoni        Won  \n",
       "1                         William Friedkin        Won  \n",
       "2                             Gene Hackman        Won  \n",
       "3                           Ernest Tidyman        Won  \n",
       "4                      Gerald B. Greenberg        Won  \n",
       "5                             Roy Scheider  Nominated  \n",
       "6                             Owen Roizman  Nominated  \n",
       "7   Theodore Soderberg\\nChristopher Newman  Nominated  \n",
       "8                      Gerald B. Greenberg  Nominated  \n",
       "9                             Gene Hackman        Won  \n",
       "10                     Gerald B. Greenberg        Won  \n",
       "11                        William Friedkin  Nominated  \n",
       "12                         Philip D'Antoni  Nominated  \n",
       "13  Christopher Newman\\nTheodore Soderberg  Nominated  \n",
       "14                         Philip D'Antoni        Won  \n",
       "15                        William Friedkin        Won  \n",
       "16                          Ernest Tidyman        Won  \n",
       "17                        Phillip D'Antoni        Won  \n",
       "18                        William Friedkin        Won  \n",
       "19                            Gene Hackman        Won  \n",
       "20                          Ernest Tidyman  Nominated  \n",
       "21                            Gene Hackman        Won  \n",
       "22                          Ernest Tidyman        Won  \n",
       "23                            Gene Hackman  Nominated  \n",
       "24                            Gene Hackman        Won  \n",
       "25                          Ernest Tidyman  Nominated  \n",
       "26                          Ernest Tidyman  Nominated  "
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "4e878e7e-1597-442a-97b3-bf3347cd4993",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_engine = ChainOfTableQueryEngine(df, llm=llm, verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "c8bbc84a-d114-4ff4-bcb2-144a595e48b8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1;3;32m> Iteration: 0\n",
      "\u001b[0m\u001b[1;3;34m> Current table:\n",
      "col : Award | Category | Nominee | Result\n",
      "row 1 : Academy Awards, 1972 | Best Picture | Phillip D'Antoni | Won\n",
      "row 2 : Academy Awards, 1972 | Best Director | William Friedkin | Won\n",
      "row 3 : Academy Awards, 1972 | Best Actor | Gene Hackman | Won\n",
      "row 4 : Academy Awards, 1972 | Best Adapted Screenplay | Ernest Tidyman | Won\n",
      "row 5 : Academy Awards, 1972 | Film Editing | Gerald B. Greenberg | Won\n",
      "row 6 : Academy Awards, 1972 | Best Supporting Actor | Roy Scheider | Nominated\n",
      "row 7 : Academy Awards, 1972 | Best Cinematography | Owen Roizman | Nominated\n",
      "row 8 : Academy Awards, 1972 | Best Sound | Theodore Soderberg\\nChristopher Newman | Nominated\n",
      "row 9 : American Cinema Editors, 1972 | Best Edited Feature Film | Gerald B. Greenberg | Nominated\n",
      "row 10 : BAFTA, 1972 | Best Actor | Gene Hackman | Won\n",
      "row 11 : BAFTA, 1972 | Best Film Editing | Gerald B. Greenberg | Won\n",
      "row 12 : BAFTA, 1972 | Best Direction | William Friedkin | Nominated\n",
      "row 13 : BAFTA, 1972 | Best Film | Philip D'Antoni | Nominated\n",
      "row 14 : BAFTA, 1972 | Best Sound Track | Christopher Newman\\nTheodore Soderberg | Nominated\n",
      "row 15 : David di Donatello Award, 1972 | Best Foreign Film | Philip D'Antoni | Won\n",
      "row 16 : Directors Guild of America, 1972 | Outstanding Directorial Achievement | William Friedkin | Won\n",
      "row 17 : Edgar Allan Poe Awards, 1972 | Best Motion Picture | Ernest Tidyman | Won\n",
      "row 18 : Golden Globe Awards, 1972 | Best Motion Picture | Phillip D'Antoni | Won\n",
      "row 19 : Golden Globe Awards, 1972 | Best Director | William Friedkin | Won\n",
      "row 20 : Golden Globe Awards, 1972 | Best Actor | Gene Hackman | Won\n",
      "row 21 : Golden Globe Awards, 1972 | Best Screenplay | Ernest Tidyman | Nominated\n",
      "row 22 : Kansas City Film Critics Circle, 1972 | Best Actor | Gene Hackman | Won\n",
      "row 23 : Kansas City Film Critics Circle, 1972 | Best Film | Ernest Tidyman | Won\n",
      "row 24 : National Society of Film Critics, 1972 | Best Actor | Gene Hackman | Nominated\n",
      "row 25 : New York Film Critics Circle, 1971 | Best Actor | Gene Hackman | Won\n",
      "row 26 : New York Film Critics Circle, 1971 | Best Film | Ernest Tidyman | Nominated\n",
      "row 27 : Writers Guild of America, 1972 | Best Drama Adaptation | Ernest Tidyman | Nominated\n",
      "\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;5;200m> New Operation + Args: f_select_row(['row 2'])\n",
      "\u001b[0m\u001b[1;3;38;5;200m> Current chain: f_select_row(['row 2']) -> \n",
      "\u001b[0m\u001b[1;3;32m> Iteration: 1\n",
      "\u001b[0m\u001b[1;3;34m> Current table:\n",
      "col : Award | Category | Nominee | Result\n",
      "row 1 : Academy Awards, 1972 | Best Director | William Friedkin | Won\n",
      "\n",
      "\n",
      "\u001b[0m> Ending operation chain.\n"
     ]
    }
   ],
   "source": [
    "response = query_engine.query(\"Who won best Director in the 1972 Academy Awards?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "21a6b620-5fcd-477f-9d3f-e6ed2a823ae0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'assistant: William Friedkin.'"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "str(response.response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2847c2b0-ad8c-4fdf-8f28-1f5f41c6c719",
   "metadata": {},
   "source": [
    "### Example: Yearly Temperature and Precipitation\n",
    "\n",
    "This table is interesting the cells for the first three rows contain two values (e.g. C/F or mm/inches).\n",
    "\n",
    "Let's see if chain-of-table can handle this question."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "b959a7e4-d2fb-4116-bf28-7e4aea990149",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "df = pd.read_csv(\"./WikiTableQuestions/csv/200-csv/42.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "feb2e453-01c8-4991-8acb-daab70241a63",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Month</th>\n",
       "      <th>Jan</th>\n",
       "      <th>Feb</th>\n",
       "      <th>Mar</th>\n",
       "      <th>Apr</th>\n",
       "      <th>May</th>\n",
       "      <th>Jun</th>\n",
       "      <th>Jul</th>\n",
       "      <th>Aug</th>\n",
       "      <th>Sep</th>\n",
       "      <th>Oct</th>\n",
       "      <th>Nov</th>\n",
       "      <th>Dec</th>\n",
       "      <th>Year</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Average high °C (°F)</td>\n",
       "      <td>17.3\\n(63.1)</td>\n",
       "      <td>19.5\\n(67.1)</td>\n",
       "      <td>22.6\\n(72.7)</td>\n",
       "      <td>25.9\\n(78.6)</td>\n",
       "      <td>27.2\\n(81)</td>\n",
       "      <td>29.3\\n(84.7)</td>\n",
       "      <td>31.8\\n(89.2)</td>\n",
       "      <td>31.4\\n(88.5)</td>\n",
       "      <td>28.9\\n(84)</td>\n",
       "      <td>25.5\\n(77.9)</td>\n",
       "      <td>21.7\\n(71.1)</td>\n",
       "      <td>19.2\\n(66.6)</td>\n",
       "      <td>24.76\\n(76.57)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Average low °C (°F)</td>\n",
       "      <td>7.9\\n(46.2)</td>\n",
       "      <td>9.4\\n(48.9)</td>\n",
       "      <td>12.5\\n(54.5)</td>\n",
       "      <td>17.6\\n(63.7)</td>\n",
       "      <td>19.2\\n(66.6)</td>\n",
       "      <td>21.6\\n(70.9)</td>\n",
       "      <td>23.8\\n(74.8)</td>\n",
       "      <td>22.5\\n(72.5)</td>\n",
       "      <td>20.7\\n(69.3)</td>\n",
       "      <td>16.5\\n(61.7)</td>\n",
       "      <td>14.1\\n(57.4)</td>\n",
       "      <td>8.5\\n(47.3)</td>\n",
       "      <td>15.94\\n(60.69)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Precipitation mm (inches)</td>\n",
       "      <td>235.9\\n(9.287)</td>\n",
       "      <td>129.2\\n(5.087)</td>\n",
       "      <td>82.8\\n(3.26)</td>\n",
       "      <td>33.6\\n(1.323)</td>\n",
       "      <td>4.7\\n(0.185)</td>\n",
       "      <td>0.2\\n(0.008)</td>\n",
       "      <td>0.0\\n(0)</td>\n",
       "      <td>0.2\\n(0.008)</td>\n",
       "      <td>3.2\\n(0.126)</td>\n",
       "      <td>58.0\\n(2.283)</td>\n",
       "      <td>107.4\\n(4.228)</td>\n",
       "      <td>214.5\\n(8.445)</td>\n",
       "      <td>857.3\\n(33.752)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Avg. precipitation days</td>\n",
       "      <td>13.9</td>\n",
       "      <td>11.4</td>\n",
       "      <td>8.6</td>\n",
       "      <td>3.6</td>\n",
       "      <td>2.4</td>\n",
       "      <td>0.1</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.1</td>\n",
       "      <td>1.8</td>\n",
       "      <td>4.9</td>\n",
       "      <td>8.0</td>\n",
       "      <td>11.8</td>\n",
       "      <td>63.7</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                       Month             Jan             Feb           Mar  \\\n",
       "0       Average high °C (°F)    17.3\\n(63.1)    19.5\\n(67.1)  22.6\\n(72.7)   \n",
       "1        Average low °C (°F)     7.9\\n(46.2)     9.4\\n(48.9)  12.5\\n(54.5)   \n",
       "2  Precipitation mm (inches)  235.9\\n(9.287)  129.2\\n(5.087)  82.8\\n(3.26)   \n",
       "3    Avg. precipitation days            13.9            11.4           8.6   \n",
       "\n",
       "             Apr           May           Jun           Jul           Aug  \\\n",
       "0   25.9\\n(78.6)    27.2\\n(81)  29.3\\n(84.7)  31.8\\n(89.2)  31.4\\n(88.5)   \n",
       "1   17.6\\n(63.7)  19.2\\n(66.6)  21.6\\n(70.9)  23.8\\n(74.8)  22.5\\n(72.5)   \n",
       "2  33.6\\n(1.323)  4.7\\n(0.185)  0.2\\n(0.008)      0.0\\n(0)  0.2\\n(0.008)   \n",
       "3            3.6           2.4           0.1           0.0           0.1   \n",
       "\n",
       "            Sep            Oct             Nov             Dec  \\\n",
       "0    28.9\\n(84)   25.5\\n(77.9)    21.7\\n(71.1)    19.2\\n(66.6)   \n",
       "1  20.7\\n(69.3)   16.5\\n(61.7)    14.1\\n(57.4)     8.5\\n(47.3)   \n",
       "2  3.2\\n(0.126)  58.0\\n(2.283)  107.4\\n(4.228)  214.5\\n(8.445)   \n",
       "3           1.8            4.9             8.0            11.8   \n",
       "\n",
       "              Year  \n",
       "0   24.76\\n(76.57)  \n",
       "1   15.94\\n(60.69)  \n",
       "2  857.3\\n(33.752)  \n",
       "3             63.7  "
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "5d9f4d93-2490-41ff-9aab-4bd251542d3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_engine = ChainOfTableQueryEngine(df, llm=llm, verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "20da9f10-42d0-451f-ad0d-38fbcd7698d0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1;3;32m> Iteration: 0\n",
      "\u001b[0m\u001b[1;3;34m> Current table:\n",
      "col : Month | Jan | Feb | Mar | Apr | May | Jun | Jul | Aug | Sep | Oct | Nov | Dec | Year\n",
      "row 1 : Average high °C (°F) | 17.3\\n(63.1) | 19.5\\n(67.1) | 22.6\\n(72.7) | 25.9\\n(78.6) | 27.2\\n(81) | 29.3\\n(84.7) | 31.8\\n(89.2) | 31.4\\n(88.5) | 28.9\\n(84) | 25.5\\n(77.9) | 21.7\\n(71.1) | 19.2\\n(66.6) | 24.76\\n(76.57)\n",
      "row 2 : Average low °C (°F) | 7.9\\n(46.2) | 9.4\\n(48.9) | 12.5\\n(54.5) | 17.6\\n(63.7) | 19.2\\n(66.6) | 21.6\\n(70.9) | 23.8\\n(74.8) | 22.5\\n(72.5) | 20.7\\n(69.3) | 16.5\\n(61.7) | 14.1\\n(57.4) | 8.5\\n(47.3) | 15.94\\n(60.69)\n",
      "row 3 : Precipitation mm (inches) | 235.9\\n(9.287) | 129.2\\n(5.087) | 82.8\\n(3.26) | 33.6\\n(1.323) | 4.7\\n(0.185) | 0.2\\n(0.008) | 0.0\\n(0) | 0.2\\n(0.008) | 3.2\\n(0.126) | 58.0\\n(2.283) | 107.4\\n(4.228) | 214.5\\n(8.445) | 857.3\\n(33.752)\n",
      "row 4 : Avg. precipitation days | 13.9 | 11.4 | 8.6 | 3.6 | 2.4 | 0.1 | 0.0 | 0.1 | 1.8 | 4.9 | 8.0 | 11.8 | 63.7\n",
      "\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;5;200m> New Operation + Args: f_select_row(['row 3'])\n",
      "\u001b[0m\u001b[1;3;38;5;200m> Current chain: f_select_row(['row 3']) -> \n",
      "\u001b[0m\u001b[1;3;32m> Iteration: 1\n",
      "\u001b[0m\u001b[1;3;34m> Current table:\n",
      "col : Month | Jan | Feb | Mar | Apr | May | Jun | Jul | Aug | Sep | Oct | Nov | Dec | Year\n",
      "row 1 : Precipitation mm (inches) | 235.9\\n(9.287) | 129.2\\n(5.087) | 82.8\\n(3.26) | 33.6\\n(1.323) | 4.7\\n(0.185) | 0.2\\n(0.008) | 0.0\\n(0) | 0.2\\n(0.008) | 3.2\\n(0.126) | 58.0\\n(2.283) | 107.4\\n(4.228) | 214.5\\n(8.445) | 857.3\\n(33.752)\n",
      "\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;5;200m> New Operation + Args: f_select_column(['Jun'])\n",
      "\u001b[0m\u001b[1;3;38;5;200m> Current chain: f_select_row(['row 3']) -> f_select_column(['Jun']) -> \n",
      "\u001b[0m\u001b[1;3;32m> Iteration: 2\n",
      "\u001b[0m\u001b[1;3;34m> Current table:\n",
      "col : Jun\n",
      "row 1 : 0.2\\n(0.008)\n",
      "\n",
      "\n",
      "\u001b[0m> Ending operation chain.\n"
     ]
    }
   ],
   "source": [
    "response = query_engine.query(\"What was the precipitation in inches during June?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "8a73c4d8-076c-4e0e-aaf5-5cb70fea472d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'assistant: 0.008 inches.'"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "str(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ffa28812-30aa-4f0c-a8a7-c20afaaf7557",
   "metadata": {},
   "source": [
    "#### Try out a Baseline\n",
    "\n",
    "As an example lets take our LLM and see if it can directly answer the question by dumping the table into the prompt! \n",
    "\n",
    "We can construct this concisely using our query pipeline syntax (you can, of course, just call the prompt/llm directly)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "1b9221c3-48d9-4668-b6ef-2b331dfa225c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.prompts import PromptTemplate\n",
    "from llama_index.query_pipeline import QueryPipeline\n",
    "\n",
    "prompt_str = \"\"\"\\\n",
    "Here's a serialized table.\n",
    "\n",
    "{serialized_table}\n",
    "\n",
    "Given this table please answer the question: {question}\n",
    "Answer: \"\"\"\n",
    "prompt = PromptTemplate(prompt_str)\n",
    "prompt_c = prompt.as_query_component(partial={\"serialized_table\": serialize_table(df)})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "806a694a-47ac-4855-93cf-570e554b06ae",
   "metadata": {},
   "source": [
    "In the response below, we see that the right row is identified, but it mistakenly identifies 0.2 as the inches instead of 0.008."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "9981850c-1e30-444b-9feb-889c886a69fd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "assistant: The precipitation in inches during June is given in row 3 under the \"Jun\" column. According to the table, it is 0.2 inches (0.008).\n"
     ]
    }
   ],
   "source": [
    "qp = QueryPipeline(chain=[prompt_c, llm])\n",
    "response = qp.run(\"What was the precipitation in inches during June?\")\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc43b093-9265-46a1-8eda-037415d529ea",
   "metadata": {},
   "source": [
    "### Example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 142,
   "id": "34637515-190e-463c-b76d-1c388de2f37d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Week</th>\n",
       "      <th>Date</th>\n",
       "      <th>TV Time</th>\n",
       "      <th>Opponent</th>\n",
       "      <th>Result</th>\n",
       "      <th>Game site</th>\n",
       "      <th>Record</th>\n",
       "      <th>Attendance</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>September 7, 1998</td>\n",
       "      <td>ABC 7:00 pm MT</td>\n",
       "      <td>New England Patriots</td>\n",
       "      <td>W 27–21</td>\n",
       "      <td>Mile High Stadium (ABC)</td>\n",
       "      <td>1–0</td>\n",
       "      <td>74,745</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2</td>\n",
       "      <td>September 13, 1998</td>\n",
       "      <td>FOX 2:00 pm MT</td>\n",
       "      <td>Dallas Cowboys</td>\n",
       "      <td>W 42–23</td>\n",
       "      <td>Mile High Stadium (FOX)</td>\n",
       "      <td>2–0</td>\n",
       "      <td>75,013</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>3</td>\n",
       "      <td>September 20, 1998</td>\n",
       "      <td>CBS 2:00 pm MT</td>\n",
       "      <td>at Oakland Raiders</td>\n",
       "      <td>W 34–17</td>\n",
       "      <td>Oakland-Alameda County Coliseum (CBS)</td>\n",
       "      <td>3–0</td>\n",
       "      <td>56,578</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>4</td>\n",
       "      <td>September 27, 1998</td>\n",
       "      <td>CBS 11:00 am MT</td>\n",
       "      <td>at Washington Redskins</td>\n",
       "      <td>W 38–16</td>\n",
       "      <td>FedEx Field (CBS)</td>\n",
       "      <td>4–0</td>\n",
       "      <td>71,880</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>5</td>\n",
       "      <td>October 4, 1998</td>\n",
       "      <td>FOX 2:00 pm MT</td>\n",
       "      <td>Philadelphia Eagles</td>\n",
       "      <td>W 41–16</td>\n",
       "      <td>Mile High Stadium (FOX)</td>\n",
       "      <td>5–0</td>\n",
       "      <td>73,218</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>6</td>\n",
       "      <td>October 11, 1998</td>\n",
       "      <td>CBS 2:00 pm MT</td>\n",
       "      <td>at Seattle Seahawks</td>\n",
       "      <td>W 21–16</td>\n",
       "      <td>Kingdome (CBS)</td>\n",
       "      <td>6–0</td>\n",
       "      <td>66,258</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>7</td>\n",
       "      <td>Bye</td>\n",
       "      <td>Bye</td>\n",
       "      <td>Bye</td>\n",
       "      <td>Bye</td>\n",
       "      <td>Bye</td>\n",
       "      <td>Bye</td>\n",
       "      <td>Bye</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>8</td>\n",
       "      <td>October 25, 1998</td>\n",
       "      <td>CBS 2:00 pm MT</td>\n",
       "      <td>Jacksonville Jaguars</td>\n",
       "      <td>W 37–24</td>\n",
       "      <td>Mile High Stadium (CBS)</td>\n",
       "      <td>7–0</td>\n",
       "      <td>75,217</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>9</td>\n",
       "      <td>November 1, 1998</td>\n",
       "      <td>CBS 11:00 am MT</td>\n",
       "      <td>at Cincinnati Bengals</td>\n",
       "      <td>W 33–26</td>\n",
       "      <td>Cinergy Field (CBS)</td>\n",
       "      <td>8–0</td>\n",
       "      <td>59,974</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>10</td>\n",
       "      <td>November 8, 1998</td>\n",
       "      <td>CBS 2:00 pm MT</td>\n",
       "      <td>San Diego Chargers</td>\n",
       "      <td>W 27–10</td>\n",
       "      <td>Mile High Stadium (CBS)</td>\n",
       "      <td>9–0</td>\n",
       "      <td>74,925</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>11</td>\n",
       "      <td>November 16, 1998</td>\n",
       "      <td>ABC 7:00 pm MT</td>\n",
       "      <td>at Kansas City Chiefs</td>\n",
       "      <td>W 30–7</td>\n",
       "      <td>Arrowhead Stadium (ABC)</td>\n",
       "      <td>10–0</td>\n",
       "      <td>78,100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>12</td>\n",
       "      <td>November 22, 1998</td>\n",
       "      <td>CBS 2:00 pm MT</td>\n",
       "      <td>Oakland Raiders</td>\n",
       "      <td>W 40–14</td>\n",
       "      <td>Mile High Stadium (CBS)</td>\n",
       "      <td>11–0</td>\n",
       "      <td>75,325</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>13</td>\n",
       "      <td>November 29, 1998</td>\n",
       "      <td>ESPN 6:15 pm MT</td>\n",
       "      <td>at San Diego Chargers</td>\n",
       "      <td>W 31–16</td>\n",
       "      <td>Qualcomm Stadium (ESPN)</td>\n",
       "      <td>12–0</td>\n",
       "      <td>66,532</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13</th>\n",
       "      <td>14</td>\n",
       "      <td>December 6, 1998</td>\n",
       "      <td>CBS 2:00 pm MT</td>\n",
       "      <td>Kansas City Chiefs</td>\n",
       "      <td>W 35–31</td>\n",
       "      <td>Mile High Stadium (CBS)</td>\n",
       "      <td>13–0</td>\n",
       "      <td>74,962</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14</th>\n",
       "      <td>15</td>\n",
       "      <td>December 13, 1998</td>\n",
       "      <td>CBS 11:00 am MT</td>\n",
       "      <td>at New York Giants</td>\n",
       "      <td>L 16–20</td>\n",
       "      <td>Giants Stadium (CBS)</td>\n",
       "      <td>13–1</td>\n",
       "      <td>72,336</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>15</th>\n",
       "      <td>16</td>\n",
       "      <td>December 21, 1998</td>\n",
       "      <td>ABC 7:00 pm MT</td>\n",
       "      <td>at Miami Dolphins</td>\n",
       "      <td>L 21–31</td>\n",
       "      <td>Pro Player Stadium (ABC)</td>\n",
       "      <td>13–2</td>\n",
       "      <td>74,363</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16</th>\n",
       "      <td>17</td>\n",
       "      <td>December 27, 1998</td>\n",
       "      <td>CBS 2:15 pm MT</td>\n",
       "      <td>Seattle Seahawks</td>\n",
       "      <td>W 28–21</td>\n",
       "      <td>Mile High Stadium (CBS)</td>\n",
       "      <td>14–2</td>\n",
       "      <td>74,057</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    Week                Date          TV Time                Opponent  \\\n",
       "0      1   September 7, 1998   ABC 7:00 pm MT    New England Patriots   \n",
       "1      2  September 13, 1998   FOX 2:00 pm MT          Dallas Cowboys   \n",
       "2      3  September 20, 1998   CBS 2:00 pm MT      at Oakland Raiders   \n",
       "3      4  September 27, 1998  CBS 11:00 am MT  at Washington Redskins   \n",
       "4      5     October 4, 1998   FOX 2:00 pm MT     Philadelphia Eagles   \n",
       "5      6    October 11, 1998   CBS 2:00 pm MT     at Seattle Seahawks   \n",
       "6      7                 Bye              Bye                     Bye   \n",
       "7      8    October 25, 1998   CBS 2:00 pm MT    Jacksonville Jaguars   \n",
       "8      9    November 1, 1998  CBS 11:00 am MT   at Cincinnati Bengals   \n",
       "9     10    November 8, 1998   CBS 2:00 pm MT      San Diego Chargers   \n",
       "10    11   November 16, 1998   ABC 7:00 pm MT   at Kansas City Chiefs   \n",
       "11    12   November 22, 1998   CBS 2:00 pm MT         Oakland Raiders   \n",
       "12    13   November 29, 1998  ESPN 6:15 pm MT   at San Diego Chargers   \n",
       "13    14    December 6, 1998   CBS 2:00 pm MT      Kansas City Chiefs   \n",
       "14    15   December 13, 1998  CBS 11:00 am MT      at New York Giants   \n",
       "15    16   December 21, 1998   ABC 7:00 pm MT       at Miami Dolphins   \n",
       "16    17   December 27, 1998   CBS 2:15 pm MT        Seattle Seahawks   \n",
       "\n",
       "     Result                              Game site Record Attendance  \n",
       "0   W 27–21                Mile High Stadium (ABC)    1–0     74,745  \n",
       "1   W 42–23                Mile High Stadium (FOX)    2–0     75,013  \n",
       "2   W 34–17  Oakland-Alameda County Coliseum (CBS)    3–0     56,578  \n",
       "3   W 38–16                      FedEx Field (CBS)    4–0     71,880  \n",
       "4   W 41–16                Mile High Stadium (FOX)    5–0     73,218  \n",
       "5   W 21–16                         Kingdome (CBS)    6–0     66,258  \n",
       "6       Bye                                    Bye    Bye        Bye  \n",
       "7   W 37–24                Mile High Stadium (CBS)    7–0     75,217  \n",
       "8   W 33–26                    Cinergy Field (CBS)    8–0     59,974  \n",
       "9   W 27–10                Mile High Stadium (CBS)    9–0     74,925  \n",
       "10   W 30–7                Arrowhead Stadium (ABC)   10–0     78,100  \n",
       "11  W 40–14                Mile High Stadium (CBS)   11–0     75,325  \n",
       "12  W 31–16                Qualcomm Stadium (ESPN)   12–0     66,532  \n",
       "13  W 35–31                Mile High Stadium (CBS)   13–0     74,962  \n",
       "14  L 16–20                   Giants Stadium (CBS)   13–1     72,336  \n",
       "15  L 21–31               Pro Player Stadium (ABC)   13–2     74,363  \n",
       "16  W 28–21                Mile High Stadium (CBS)   14–2     74,057  "
      ]
     },
     "execution_count": 142,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "\n",
    "df = pd.read_csv(\"./WikiTableQuestions/csv/203-csv/114.csv\")\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 144,
   "id": "65dbe340-1fd9-47b3-9ce0-023aa13416b1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1;3;32m> Iteration: 0\n",
      "\u001b[0m\u001b[1;3;34m> Current table:\n",
      "col : Week | Date | TV Time | Opponent | Result | Game site | Record | Attendance\n",
      "row 1 : 1 | September 7, 1998 | ABC 7:00 pm MT | New England Patriots | W 27–21 | Mile High Stadium (ABC) | 1–0 | 74,745\n",
      "row 2 : 2 | September 13, 1998 | FOX 2:00 pm MT | Dallas Cowboys | W 42–23 | Mile High Stadium (FOX) | 2–0 | 75,013\n",
      "row 3 : 3 | September 20, 1998 | CBS 2:00 pm MT | at Oakland Raiders | W 34–17 | Oakland-Alameda County Coliseum (CBS) | 3–0 | 56,578\n",
      "row 4 : 4 | September 27, 1998 | CBS 11:00 am MT | at Washington Redskins | W 38–16 | FedEx Field (CBS) | 4–0 | 71,880\n",
      "row 5 : 5 | October 4, 1998 | FOX 2:00 pm MT | Philadelphia Eagles | W 41–16 | Mile High Stadium (FOX) | 5–0 | 73,218\n",
      "row 6 : 6 | October 11, 1998 | CBS 2:00 pm MT | at Seattle Seahawks | W 21–16 | Kingdome (CBS) | 6–0 | 66,258\n",
      "row 7 : 7 | Bye | Bye | Bye | Bye | Bye | Bye | Bye\n",
      "row 8 : 8 | October 25, 1998 | CBS 2:00 pm MT | Jacksonville Jaguars | W 37–24 | Mile High Stadium (CBS) | 7–0 | 75,217\n",
      "row 9 : 9 | November 1, 1998 | CBS 11:00 am MT | at Cincinnati Bengals | W 33–26 | Cinergy Field (CBS) | 8–0 | 59,974\n",
      "row 10 : 10 | November 8, 1998 | CBS 2:00 pm MT | San Diego Chargers | W 27–10 | Mile High Stadium (CBS) | 9–0 | 74,925\n",
      "row 11 : 11 | November 16, 1998 | ABC 7:00 pm MT | at Kansas City Chiefs | W 30–7 | Arrowhead Stadium (ABC) | 10–0 | 78,100\n",
      "row 12 : 12 | November 22, 1998 | CBS 2:00 pm MT | Oakland Raiders | W 40–14 | Mile High Stadium (CBS) | 11–0 | 75,325\n",
      "row 13 : 13 | November 29, 1998 | ESPN 6:15 pm MT | at San Diego Chargers | W 31–16 | Qualcomm Stadium (ESPN) | 12–0 | 66,532\n",
      "row 14 : 14 | December 6, 1998 | CBS 2:00 pm MT | Kansas City Chiefs | W 35–31 | Mile High Stadium (CBS) | 13–0 | 74,962\n",
      "row 15 : 15 | December 13, 1998 | CBS 11:00 am MT | at New York Giants | L 16–20 | Giants Stadium (CBS) | 13–1 | 72,336\n",
      "row 16 : 16 | December 21, 1998 | ABC 7:00 pm MT | at Miami Dolphins | L 21–31 | Pro Player Stadium (ABC) | 13–2 | 74,363\n",
      "row 17 : 17 | December 27, 1998 | CBS 2:15 pm MT | Seattle Seahawks | W 28–21 | Mile High Stadium (CBS) | 14–2 | 74,057\n",
      "\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;5;200m> New Operation + Args: f_select_column(['TV Time', 'Attendance'])\n",
      "\u001b[0m\u001b[1;3;38;5;200m> Current chain: f_select_column(['TV Time', 'Attendance']) -> \n",
      "\u001b[0m\u001b[1;3;32m> Iteration: 1\n",
      "\u001b[0m\u001b[1;3;34m> Current table:\n",
      "col : TV Time | Attendance\n",
      "row 1 : ABC 7:00 pm MT | 74,745\n",
      "row 2 : FOX 2:00 pm MT | 75,013\n",
      "row 3 : CBS 2:00 pm MT | 56,578\n",
      "row 4 : CBS 11:00 am MT | 71,880\n",
      "row 5 : FOX 2:00 pm MT | 73,218\n",
      "row 6 : CBS 2:00 pm MT | 66,258\n",
      "row 7 : Bye | Bye\n",
      "row 8 : CBS 2:00 pm MT | 75,217\n",
      "row 9 : CBS 11:00 am MT | 59,974\n",
      "row 10 : CBS 2:00 pm MT | 74,925\n",
      "row 11 : ABC 7:00 pm MT | 78,100\n",
      "row 12 : CBS 2:00 pm MT | 75,325\n",
      "row 13 : ESPN 6:15 pm MT | 66,532\n",
      "row 14 : CBS 2:00 pm MT | 74,962\n",
      "row 15 : CBS 11:00 am MT | 72,336\n",
      "row 16 : ABC 7:00 pm MT | 74,363\n",
      "row 17 : CBS 2:15 pm MT | 74,057\n",
      "\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;5;200m> New Operation + Args: f_select_row(['row 11'])\n",
      "\u001b[0m\u001b[1;3;38;5;200m> Current chain: f_select_column(['TV Time', 'Attendance']) -> f_select_row(['row 11']) -> \n",
      "\u001b[0m\u001b[1;3;32m> Iteration: 2\n",
      "\u001b[0m\u001b[1;3;34m> Current table:\n",
      "col : TV Time | Attendance\n",
      "row 1 : ABC 7:00 pm MT | 78,100\n",
      "\n",
      "\n",
      "\u001b[0m> Ending operation chain.\n"
     ]
    }
   ],
   "source": [
    "query_engine = ChainOfTableQueryEngine(df, llm=llm, verbose=True)\n",
    "response = query_engine.query(\"Which televised ABC game had the greatest attendance?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 145,
   "id": "f6dd33f5-ac5c-4075-b94e-a8d0da25c231",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "assistant: The answer is: ABC 7:00 pm MT with an attendance of 78,100.\n"
     ]
    }
   ],
   "source": [
    "print(str(response))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea52a9ac-4152-44de-a2cd-30fb7f0cabcc",
   "metadata": {},
   "source": [
    "#### Baseline\n",
    "\n",
    "Once again, we consider a simple QA prompt baseline and get the wrong answer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 140,
   "id": "95122085-49c9-4b61-8a01-d529beba2dcf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "assistant: According to the table, there are two games that were televised on ABC:\n",
      "\n",
      "1. Week 1: September 7, 1998, against the New England Patriots with an attendance of 74,745.\n",
      "2. Week 16: December 21, 1998, against the Miami Dolphins with an attendance of 74,363.\n",
      "\n",
      "The game with the greatest attendance among the ABC televised games is the Week 1 game against the New England Patriots, with an attendance of 74,745.\n"
     ]
    }
   ],
   "source": [
    "from llama_index.prompts import PromptTemplate\n",
    "from llama_index.query_pipeline import QueryPipeline\n",
    "\n",
    "prompt_str = \"\"\"\\\n",
    "Here's a serialized table.\n",
    "\n",
    "{serialized_table}\n",
    "\n",
    "Given this table please answer the question: {question}\n",
    "Answer: \"\"\"\n",
    "prompt = PromptTemplate(prompt_str)\n",
    "prompt_c = prompt.as_query_component(partial={\"serialized_table\": serialize_table(df)})\n",
    "qp = QueryPipeline(chain=[prompt_c, llm])\n",
    "response = qp.run(\"Which televised ABC game had the greatest attendance?\")\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1c705df-30b8-4888-b958-32432be8d774",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
